mod error;

use self::error::Error;
use crate::{
    courier::{
        codec::{ChannelEndCodec, DefaultSerder, FromSliceDeserializer, ToWriterSerializer},
        ConnectionSource, RawChannelEnd, TcpConnectionStream,
    },
    utils::SpawnedTasksClearGuard,
    ClientLocalRequestId, RpcCallRequest, RpcCancelRequest, RpcRequest, RpcResponse,
};
use fnv::FnvHashMap;
use futures::{
    ready,
    stream::{AbortHandle, AbortRegistration, Abortable},
    FutureExt, Sink, Stream,
};
use log::error;
use pin_project::pin_project;
use serde::{de::DeserializeOwned, Serialize};
use std::{
    collections::hash_map::Entry,
    fmt::Debug,
    future::Future,
    net::SocketAddr,
    panic::AssertUnwindSafe,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
};
use tokio::{
    io::{AsyncRead, AsyncWrite},
    sync::mpsc,
};

const REQ_ONESHOT_RECV_LIMIT: usize = 100;
const RESP_ONESHOT_SEND_LIMIT: usize = 100;

pub trait ServeRpc<I>: Sync + Send + 'static {
    type Req: Send;
    type Resp: Send;
    fn dispatch_call(&self, req: Self::Req, from: I) -> impl Future<Output = Self::Resp> + Send;
}

pub struct Server<S, C = TcpConnectionStream, SD = DefaultSerder> {
    serve: Arc<S>,
    conn_source: C,
    serder: SD,
}

impl<S, C, SD> Server<S, C, SD> {
    pub fn with_conn_source_and_serder(serve: Arc<S>, conn_source: C, serder: SD) -> Self {
        Self {
            serve,
            conn_source,
            serder,
        }
    }
}
impl<S, C> Server<S, C> {
    #[cfg(any(feature = "msgpack"))]
    pub fn with_default_serder(serve: Arc<S>, conn_source: C) -> Self {
        Self::with_conn_source_and_serder(serve, conn_source, DefaultSerder::default())
    }
}
impl<S> Server<S> {
    #[cfg(any(feature = "msgpack"))]
    pub async fn new_via_tcp(serve: Arc<S>, addr: SocketAddr) -> std::io::Result<Self> {
        Ok(Self::with_default_serder(
            serve,
            TcpConnectionStream::new(addr).await?,
        ))
    }
}

impl<S, C, SD> Server<S, C, SD>
where
    S: ServeRpc<C::OtherEndInfo>,
    S::Req: DeserializeOwned,
    S::Resp: Serialize,
    C: ConnectionSource,
    C::Conn: Send + 'static,
    C::OtherEndInfo: Clone + Send + 'static,
    SD: FromSliceDeserializer + ToWriterSerializer + Clone + Send + 'static,
    <SD as FromSliceDeserializer>::Err: Send + Debug,
    <SD as ToWriterSerializer>::Err: Send + Debug,
{
    pub async fn serve(self) -> Result<(), C::Err> {
        let mut spawned_tasks = SpawnedTasksClearGuard::new();
        loop {
            match self.conn_source.get_one().await {
                Ok((conn, info)) => {
                    let raw_end = RawChannelEnd::new(
                        conn,
                        ChannelEndCodec::<RpcResponse<S::Resp>, RpcRequest<S::Req>, _>::new(
                            self.serder.clone(),
                        ),
                    );
                    let (tx, rx) = mpsc::unbounded_channel();
                    let serving = ServingEnd {
                        who_is_that: info,
                        raw_end,
                        serve: Arc::clone(&self.serve),
                        send_chan_rx: rx,
                        send_chan_tx: tx,
                        ongoing_reqs: Default::default(),
                        spawned_tasks_guard: SpawnedTasksClearGuard::new(),
                    };
                    spawned_tasks.push(
                        tokio::spawn(async move {
                            match AssertUnwindSafe(serving).catch_unwind().await {
                                Err(_e_panic) => {
                                    error!("unexpected panic from a client serving task!")
                                }
                                Ok(Err(e)) => {
                                    error!("A client serving task exit with error: {e:?}")
                                }
                                _ => {}
                            }
                        })
                        .abort_handle(),
                    );
                }
                Err(e) => break Err(e),
            }
        }
    }
}

/// The serving end of **one** RPC channel.
#[pin_project]
struct ServingEnd<C, SD, S: ServeRpc<A>, A> {
    who_is_that: A,
    #[pin]
    raw_end: RawChannelEnd<C, SD, RpcResponse<S::Resp>, RpcRequest<S::Req>>,
    serve: Arc<S>,
    send_chan_rx: mpsc::UnboundedReceiver<RpcResponse<S::Resp>>,
    ongoing_reqs: FnvHashMap<ClientLocalRequestId, AbortHandle>,

    spawned_tasks_guard: SpawnedTasksClearGuard,
    send_chan_tx: mpsc::UnboundedSender<RpcResponse<S::Resp>>,
}

enum PollReceiveDone {
    UnderlyingChannelClosed,
    OneshotReceiveLimitReached,
}

enum PollSendDone {
    OneshotSendLimitReached,
}

impl<C, SD, S, A> ServingEnd<C, SD, S, A>
where
    C: AsyncWrite + AsyncRead,
    SD: FromSliceDeserializer + ToWriterSerializer,
    S: ServeRpc<A>,
    S::Resp: Serialize,
    S::Req: DeserializeOwned,
    A: Clone + Send + 'static,
{
    fn handle_req(mut self: Pin<&mut Self>, req: RpcRequest<S::Req>) {
        match req {
            RpcRequest::Call(RpcCallRequest { id, req }) => {
                let Entry::Vacant(entry) = self.as_mut().project().ongoing_reqs.entry(id) else {
                    return;
                };
                let (abort_handle, abort_regist) = AbortHandle::new_pair();
                entry.insert(abort_handle);

                let rs = RequestServing {
                    who_is_that: self.who_is_that.clone(),
                    req_id: id,
                    req,
                    serve: Arc::clone(&self.serve),
                    result_chan: self.send_chan_tx.clone(),
                    abort_regist,
                };
                self.project()
                    .spawned_tasks_guard
                    .push(tokio::spawn(rs.overwatch_serve()).abort_handle());
            }
            RpcRequest::Cancel(RpcCancelRequest { id }) => {
                if let Some(abort_handle) = self.project().ongoing_reqs.remove(&id) {
                    abort_handle.abort()
                }
            }
        }
    }

    /// Try to receive:
    /// 1. requests from the underlying data channel
    fn poll_receive(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<
        Result<
            PollReceiveDone,
            Error<<SD as ToWriterSerializer>::Err, <SD as FromSliceDeserializer>::Err>,
        >,
    > {
        for _ in 0..REQ_ONESHOT_RECV_LIMIT {
            match ready!(self.as_mut().project().raw_end.as_mut().poll_next(cx)) {
                Some(Ok(req)) => self.as_mut().handle_req(req),
                Some(Err(e)) => return Poll::Ready(Err(e.into())),
                None => return Poll::Ready(Ok(PollReceiveDone::UnderlyingChannelClosed)),
            }
        }
        Poll::Ready(Ok(PollReceiveDone::OneshotReceiveLimitReached))
    }

    /// Try to send:
    /// 1. responses from the response queue
    fn poll_send(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<
        Result<
            PollSendDone,
            Error<<SD as ToWriterSerializer>::Err, <SD as FromSliceDeserializer>::Err>,
        >,
    > {
        for _ in 0..RESP_ONESHOT_SEND_LIMIT {
            ready!(self.as_mut().project().raw_end.poll_ready(cx))?;

            match self.as_mut().project().send_chan_rx.poll_recv(cx) {
                Poll::Ready(Some(resp)) => {
                    if self
                        .as_mut()
                        .project()
                        .ongoing_reqs
                        .remove(&resp.id())
                        .is_some()
                    {
                        self.as_mut().project().raw_end.start_send(resp)?;
                    }
                }
                Poll::Ready(None) => unreachable!("Unexpected state: server send chan closed."),
                Poll::Pending => {
                    ready!(self.as_mut().project().raw_end.poll_flush(cx))?;
                    return Poll::Pending;
                }
            }
        }
        Poll::Ready(Ok(PollSendDone::OneshotSendLimitReached))
    }
}

impl<C, SD, S, A> Future for ServingEnd<C, SD, S, A>
where
    C: AsyncWrite + AsyncRead,
    SD: FromSliceDeserializer + ToWriterSerializer,
    S: ServeRpc<A>,
    S::Resp: Serialize,
    S::Req: DeserializeOwned,
    A: Clone + Send + 'static,
{
    type Output =
        Result<(), Error<<SD as ToWriterSerializer>::Err, <SD as FromSliceDeserializer>::Err>>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        loop {
            match (
                self.as_mut().poll_receive(cx)?,
                self.as_mut().poll_send(cx)?,
            ) {
                (Poll::Ready(PollReceiveDone::UnderlyingChannelClosed), poll_send) => {
                    if self.ongoing_reqs.is_empty() {
                        return Poll::Ready(Ok(()));
                    }
                    let _ = ready!(poll_send);
                }
                (Poll::Pending, Poll::Pending) => return Poll::Pending,
                _ => {}
            }
        }
    }
}

/// The serving task for **one** request.
struct RequestServing<A, S: ServeRpc<A>> {
    who_is_that: A,
    req_id: ClientLocalRequestId,
    req: S::Req,
    serve: Arc<S>,
    result_chan: mpsc::UnboundedSender<RpcResponse<S::Resp>>,
    abort_regist: AbortRegistration,
}

impl<A, S> RequestServing<A, S>
where
    S: ServeRpc<A>,
{
    async fn overwatch_serve(self) {
        let Self {
            req_id,
            req,
            who_is_that,
            serve,
            result_chan,
            abort_regist,
        } = self;
        let f = Abortable::new(
            AssertUnwindSafe(serve.dispatch_call(req, who_is_that)).catch_unwind(),
            abort_regist,
        );
        match f.await {
            Ok(Ok(resp)) => {
                let _ = result_chan.send(RpcResponse::Done { req_id, resp });
            }
            Ok(Err(_)) => {
                // the task panicked
                let _ = result_chan.send(RpcResponse::Errored { req_id });
            }
            Err(_) => {
                // the task is aborted via the AbortHandle (i.e. cancelled)
            }
        }
    }
}
